import pandas as pd
from itertools import combinations  # 必须添加的导入
from yi_util import *

'''============================================ 数据计算 BEGIN ============================================'''



''' 计算涨停率 '''
def cal_zt_ratio(all_item_df):
    tmp_df = all_item_df[~all_item_df.isin(['--']).any(axis=1)]
    tmp_df = tmp_df[tmp_df['quote_rate']> 7 ].reset_index(drop=True)
    industry_df = pd.DataFrame(tmp_df['concept_name'].value_counts())
    industry_df = industry_df.reset_index()
    industry_df.columns = ['concept_name', 'pass_num']
    tmp_df = all_item_df[['concept_name','item_count']].drop_duplicates().reset_index(drop=True)
    df_merged = tmp_df.merge(industry_df, on=['concept_name'])
    df_merged['zt_ratio'] = 100*df_merged['pass_num']/df_merged['item_count']
    df_merged = df_merged.round(2).sort_values(by='zt_ratio', ascending=False)
    return df_merged


"""
    计算两个概念的交叉比例。
    
    参数：
    - df: 包含 'concept' 和 'name' 列的 DataFrame。
    - cpt1: 概念 1 的名称。
    - cpt2: 概念 2 的名称。
    
    返回：
    - ratio1: 概念 1 中属于概念 2 的比例。
    - ratio2: 概念 2 中属于概念 1 的比例。
    """
def cal_gn_inter_ratio(df, cpt1, cpt2):
    # 获取概念1和概念2的项目
    items1 = df[df['concept_name'] == cpt1]['name'].unique()
    items2 = df[df['concept_name'] == cpt2]['name'].unique()
    # 计算交集
    inter = set(items1) & set(items2)
    inter_count = len(inter)
    # 计算各自概念的数量
    count1 = len(items1)
    count2 = len(items2)
    # 计算交叉比例
    ratio1 = inter_count / count1 if count1 != 0 else 0
    ratio2 = inter_count / count2 if count2 != 0 else 0
    return round(ratio1,2), round(ratio2,2)


'''  纵向查看大单数据: 多单连续流入多日 '''
def cal_xl_inflow(today_str):
    f_name = 'input/high_'+str(today_str)+'.csv'
    f_name = 'input/flow20_'+str(today_str)+'.csv'

    origin_df = pd.read_csv(f_name)
    flow_df = pd.read_csv(f_name)
    
    origin_df,flow_df = get_origin_flow(today_str)
    origin_df = get_code_last_data(origin_df).drop_duplicates(subset=['code'])
    origin_df = origin_df[['code','quote_rate','r3','r10']]
    # print(origin_df.columns)
    
    df = flow_df
    df['date'] = pd.to_datetime(df['date'])
    latest_date = df['date'].max()
    start_date = latest_date - pd.Timedelta(days=3)
    recent_data = df[(df['date'] >= start_date) & (df['date'] <= latest_date)]
    # 按code分组，筛选出所有xl_yoy值都大于0的code
    grouped = recent_data.groupby('code')
    valid_codes = grouped.filter(lambda x: (x['xl_yoy'] > 0).all())
    # 从valid_codes中提取每个code的近4天的xl_yoy值
    result = valid_codes.pivot(index='code', columns='date', values='xl_yoy')
    # 重命名列，使其符合要求
    result.columns = [f'xl_yoy{i+1}' for i in range(len(result.columns))]
    # 重置索引，使code成为一列
    result = result.reset_index()
    result = pd.merge(result, origin_df, on='code')
    result = result[result['quote_rate']>4].reset_index(drop=True)
    return result

# xl_df = test_xl()
# xl_df

'''  计算5分钟大单 '''
def cal_5minute_flow(flow_df):
     # 净特大单 = 主被动买入特大单 - 主被动卖出特大单
    flow_df['xl'] = (flow_df['act_buy_xl'] + flow_df['pas_buy_xl']) - (flow_df['act_sell_xl'] + flow_df['pas_sell_xl'])
    flow_df = flow_df.applymap(divide_by_100)
    flow_df = flow_df[['trade_date', 'code', 'xl']]
    flow_df = flow_df.sort_values(by=['trade_date'], ascending=True)
    flow_df = flow_df[['trade_date','code', 'xl']]
    return flow_df

''' 计算各因子与第二日涨幅raise的相关性'''
def cal_corr(df):
    df['date'] = pd.to_datetime(df['date'])
    df_prev = df.groupby('code').shift(1).add_suffix('_prev')
    df_combined = pd.concat([df, df_prev], axis=1)
    columns_to_exclude = ['date', 'code', 'date_prev', 'code_prev']
    columns_to_include = df_combined.columns.difference(columns_to_exclude)
    df_combined = df_combined[columns_to_include].dropna()
    corr_result= df_combined.corr()['raise'].sort_values().round(2)
    return df_combined, corr_result

"""
    计算概念板块间的交叉比例
    
    参数：
    gn_df : DataFrame 个股列表数据，包含列 ['concept_name', 'name', 'quote_rate']
    zt_df : DataFrame 概念板数据，包含列 ['concept_name']
    min_quote_rate : int 最低引用率阈值，默认7
    format_percentage : bool 是否格式化为百分比字符串
    
    返回：
    DataFrame 包含列 ['concept_a', 'concept_b', 'common_count', 'cross_ratio']
"""

def cal_concept_cross_ratio(gn_df, zt_df, min_quote_rate=7, format_percentage=False):
    required_cols = ['concept_name', 'name', 'quote_rate']
    if not all(col in gn_df.columns for col in required_cols):
        raise ValueError(f"输入数据需包含列：{required_cols}")
    valid_concepts = zt_df['concept_name'].unique()
    filtered = gn_df[
        (gn_df['concept_name'].isin(valid_concepts)) &
        (gn_df['quote_rate'] > min_quote_rate)
    ].drop_duplicates(['concept_name', 'name'])
    concept_dict = filtered.groupby('concept_name')['name'].apply(set).to_dict()
    concept_pairs = list(combinations(concept_dict.keys(), 2))
    results = []
    for a, b in concept_pairs:
        common = len(concept_dict[a] & concept_dict[b])
        total = len(concept_dict[a])
        ratio = common / total if total > 0 else 0
        
        results.append({
            'concept_a': a,
            'concept_b': b,
            'common_count': common,
            'cross_ratio': ratio
        })
    result_df = pd.DataFrame(results)
    if len(result_df) == 0:
        return result_df
    result_df = result_df[result_df['cross_ratio'] > 0.2]
    if format_percentage:
        result_df['cross_ratio'] = result_df['cross_ratio'].apply(lambda x: f"{x:.2%}")
    return sort_df(result_df,'cross_ratio')
# output = cal_concept_cross_ratio(gn_list, zt_ratio_df, format_percentage=True)
    


'''============================================ 数据计算 END ============================================'''

